# compare with ~/xchg/Match/cem/a1lalonde_nocaliper.R, except with estimand=ATT

#the truth for lalonde, regression model

#compare with lapo:~/xchg/Match/GenMatch/mc1/lalonde3matching1.Rout

library(Matching)
data(lalonde)
source("AutoCluster3.R")

sims <- 1000
estimand <- "ATT"
obs <- nrow(lalonde)

#setDATA seed
set.seed(2810192)

attach(lalonde)
propensity  <- glm(treat~ I(age^2) + I(educ^2) + black +
                   hisp + married + nodegr + I(re74^2) + I(re75^2) +
                   u74 + u75, family=binomial, data=lalonde)

X <- cbind(rep(1, length(treat)),
           propensity$linear.pred,
           I(log(age)^2),
           I(log(educ)^2),
           I(log(re74+0.01)^2),
           I(log(re75+0.01)^2))

propensity.coeffs <- propensity$coeff
propensity.coeffs <- as.matrix(c(
                                 1+00,  #(Intercept)        
                                 .5,    #linear.pred
                                 .01, #age
                                 -.3,  #educ
                                 -0.01, #I(re74^2)          
                                 0.01   #I(re75^2)          
                                 ))

mu = X %*% propensity.coeffs
Tr.pred <- exp(mu)/(1+exp(mu))
#print(summary(Tr.pred))

TreatmentEffect <- 1000
TreatmentReal <- matrix(nrow=obs, ncol=1)

#return objects
raw <- matrix(nrow=sims)
correct.est.bias <- matrix(nrow=sims)
correct.est.nobias <- matrix(nrow=sims)
est.bias <- matrix(nrow=sims)
est.nobias <- matrix(nrow=sims)
iv.est.bias <- matrix(nrow=sims)
iv.est.nobias <- matrix(nrow=sims)
ps.est.bias <- matrix(nrow=sims)
ps.est.nobias <- matrix(nrow=sims)
genmatch.est.bias <- matrix(nrow=sims)
genmatch.est.nobias <- matrix(nrow=sims)


results <- function(s)
  {
    cat("RAW bias:", mean(raw[1:s]), "\n")
    cat("RAW RMSE:", sqrt(mean(raw[1:s]^2)), "\n")    
    
    cat("\n")
    
    cat("MH estimate with BIAS correction:", mean(est.bias[1:s]), "\n")
    cat("MH estimate with BIAS correction RMSE:", sqrt(mean(est.bias[1:s]^2)), "\n")
    
    cat("\n")
    
    cat("IV estimate with BIAS correction:", mean(iv.est.bias[1:s]), "\n")
    cat("IV estimate with BIAS correction RMSE:", sqrt(mean(iv.est.bias[1:s]^2)), "\n")
    
    cat("\n")
    
    cat("ps estimate with BIAS correction:", mean(ps.est.bias[1:s]), "\n")
    cat("ps estimate with BIAS correction RMSE:", sqrt(mean(ps.est.bias[1:s]^2)), "\n")
    
    cat("\n")
    
    cat("genmatch estimate with BIAS correction:", mean(genmatch.est.bias[1:s]), "\n")
    cat("genmatch estimate with BIAS correction RMSE:", sqrt(mean(genmatch.est.bias[1:s]^2)), "\n")
    
    cat("\n")
    
    cat("MH estimate w/out bias correction:", mean(est.nobias[1:s]), "\n")
    cat("MH estimate w/out bias correction RMSE:", sqrt(mean(est.nobias[1:s]^2)), "\n")
    
    cat("\n")
    
    cat("IV estimate w/out bias correction:", mean(iv.est.nobias[1:s]), "\n")
    cat("IV estimate w/out bias correction RMSE:", sqrt(mean(iv.est.nobias[1:s]^2)), "\n")
    
    cat("\n")
    
    cat("ps estimate w/out bias correction:", mean(ps.est.nobias[1:s]), "\n")
    cat("ps estimate w/out bias correction RMSE:", sqrt(mean(ps.est.nobias[1:s]^2)), "\n")
    
    cat("\n")
    
    cat("genmatch estimate w/out bias correction:", mean(genmatch.est.nobias[1:s]), "\n")
    cat("genmatch estimate w/out bias correction RMSE:", sqrt(mean(genmatch.est.nobias[1:s]^2)), "\n")
    
    cat("\n")
    
  }

count <- 0
for (s in 1:sims)
  {

    cat("\nS:", s, "\n")
    count <-  count+1
    
    for(i in 1:obs)
      {
        TreatmentReal[i] = sample(0:1, 1, prob=c(1-Tr.pred[i],Tr.pred[i]))
      }

    Y <- I(TreatmentEffect*TreatmentReal) + .1*exp(.7*log(re74+0.01) + .7*log(re75+0.01)) + rnorm(obs, 0, 10)

    raw[s] = mean(Y[TreatmentReal==1])-mean(Y[TreatmentReal==0]) - TreatmentEffect
    cat("raw[",s,"]:", raw[s], "\n")
    cat("raw old[",s,"]:", (mean(Y[TreatmentReal==1])-mean(Y[TreatmentReal==0])) , "\n")    

    glm1  <- glm(TreatmentReal~X-1,
                 family=binomial)
    #print(summary(glm1))
    Xtmp = cbind(age, educ, black, hisp, married, nodegr, re74, re75, u74, u75)
    
    X  <- glm1$linear.pred
    Y  <- Y
    
    rr  <- Match(Y=Y,Tr=TreatmentReal,X=X,M=1, estimand=estimand, BiasAdj=TRUE,
                 Z=Xtmp);

    correct.est.bias[s] = rr$est-TreatmentEffect
    correct.est.nobias[s] = rr$est.noadj-TreatmentEffect

    cat("correct.est.nobias[",s,"]:",correct.est.nobias[s],"\n")
    cat("correct.est.bias[",s,"]:",correct.est.bias[s],"\n")    

    m1 <- Match(Y=Y, Tr=TreatmentReal, X=Xtmp, BiasAdj=TRUE,
                Weight=2, Z=Xtmp)
    est.bias[s] = m1$est-TreatmentEffect
    est.nobias[s] = m1$est.noadj-TreatmentEffect

    cat("mh.est.nobias[",s,"]:",est.nobias[s],"\n")
    cat("mh.est.bias[",s,"]:",est.bias[s],"\n")

    miv <- Match(Y=Y, Tr=TreatmentReal, X=Xtmp, estimand=estimand, BiasAdj=TRUE,
                Weight=1, Z=Xtmp)

    iv.est.bias[s] = miv$est-TreatmentEffect
    iv.est.nobias[s] = miv$est.noadj-TreatmentEffect

    cat("iv.est.nobias[",s,"]:",iv.est.nobias[s],"\n")
    cat("iv.est.bias[",s,"]:",iv.est.bias[s],"\n")                        

    glm2  <- glm(TreatmentReal~age + educ + black +
                 hisp + married + nodegr + I(re74)  + I(re75),  
                 family=binomial)
    
    X  <- glm2$linear.pred
    mps  <- Match(Y=Y,Tr=TreatmentReal,X=X,M=1, estimand=estimand,
                  BiasAdj=TRUE, Z=Xtmp);

    ps.est.bias[s] = mps$est-TreatmentEffect
    ps.est.nobias[s] = mps$est.noadj-TreatmentEffect

    cat("ps.est.nobias[",s,"]:",ps.est.nobias[s],"\n")
    cat("ps.est.bias[",s,"]:",ps.est.bias[s],"\n")

    X2 = Xtmp
    for (i in 1:ncol(Xtmp))
      {
        lm1 = lm(Xtmp[,i]~glm2$linear.pred)
        X2[,i] = lm1$residual
      }
    X2 = cbind(glm2$linear.pred, X2)
    #genmatch
    cl <- NCPUS()
    gm1 <- GenMatch(Tr=TreatmentReal, X=X2, BalanceMatrix=Xtmp, estimand=estimand,
                    starting.values=c(1000,rep(0, ncol(X2)-1)),
                    pop.size=1000, cl=cl)
    stopCluster(cl)
    genmatch <- Match(Y=Y, Tr=TreatmentReal, X=X2, Z=Xtmp, BiasAdj=TRUE,
                  estimand=estimand, Weight.matrix=gm1)
    genmatch.est.bias[s] = genmatch$est-TreatmentEffect
    genmatch.est.nobias[s] = genmatch$est.noadj-TreatmentEffect
    cat("genmatch.est.nobias[",s,"]:",genmatch.est.nobias[s],"\n")
    cat("genmatch.est.bias[",s,"]:",genmatch.est.bias[s],"\n")    
    
    if (count>1)
      {
        results(s)
#        count <-  0
      }
  } #end of sims

results(s)



